
#include "seeta_alignment.h"

#include <iostream>
using namespace std;

SeetaAlignment::SeetaAlignment(const string model_path_face_detection,
        const string model_path_face_alignment)
    : model_path_fd(model_path_face_detection), model_path_fa(model_path_face_alignment)
{
    // Initialize face detection model
    //seeta::FaceDetection detector("../../../FaceDetection/model/seeta_fd_frontal_v1.0.bin");
    detector = new seeta::FaceDetection(model_path_fd.c_str());
    detector->SetMinFaceSize(40);
    detector->SetScoreThresh(2.f);
    detector->SetImagePyramidScaleFactor(0.8f);
    detector->SetWindowStep(4, 4);
    
    // Initialize face alignment model 
    //seeta::FaceAlignment point_detector((MODEL_DIR + "seeta_fa_v1.1.bin").c_str());
    point_detector = new seeta::FaceAlignment(model_path_fa.c_str());
    
    cout << "SeetaAlignment init finished." << endl;
}

int SeetaAlignment::alignTarget(const string &in_img, string &res_img)
{
    //cout << PATH_MAX << endl;
    bool isEnvOK = true;
    char abs_path_buffer[200];
    getcwd(abs_path_buffer , 200);
    cout << "Current dir: " << abs_path_buffer << endl;
    
    string FaceDetectionModel = model_path_fd;
    if(realpath(FaceDetectionModel.c_str(), abs_path_buffer)) {
        cout << "FaceDetectionModelPath: " << abs_path_buffer << endl;
    } else {
        cout << "No such file: " << FaceDetectionModel << endl;
        isEnvOK = false;
    }
    
    if(!isEnvOK) {
        cout << "Environment error: Check your current directory.\nQuit ..." << endl;
        return -1;
    }

    char buffer_image_path[PATH_MAX];
    strcpy(buffer_image_path, in_img.c_str());
    
    //load image
    IplImage *img_grayscale = NULL;
    //img_grayscale = cvLoadImage((DATA_DIR + "image_0001.png").c_str(), 0);
    img_grayscale = cvLoadImage(buffer_image_path, 0);
    if (img_grayscale == NULL)
    {
        return -1;
    }
    
    //IplImage *img_color = cvLoadImage((DATA_DIR + "image_0001.png").c_str(), 1);
    IplImage *img_color = cvLoadImage(buffer_image_path, 1);
    int pts_num = 5;
    int im_width = img_grayscale->width;
    int im_height = img_grayscale->height;
    unsigned char* data = new unsigned char[im_width * im_height];
    unsigned char* data_ptr = data;
    unsigned char* image_data_ptr = (unsigned char*)img_grayscale->imageData;
    int h = 0;
    for (h = 0; h < im_height; h++) {
        memcpy(data_ptr, image_data_ptr, im_width);
        data_ptr += im_width;
        image_data_ptr += img_grayscale->widthStep;
    }
    
    seeta::ImageData image_data;
    image_data.data = data;
    image_data.width = im_width;
    image_data.height = im_height;
    image_data.num_channels = 1;
    
    // Detect faces
    vector<seeta::FaceInfo> faces = detector->Detect(image_data);
    int32_t face_num = static_cast<int32_t>(faces.size());
    
    if (face_num == 0)
    {
        delete[]data;
        cvReleaseImage(&img_grayscale);
        cvReleaseImage(&img_color);
        return 0;
    }
    
    // Detect 5 facial landmarks
    seeta::FacialLandmark points[5];
    point_detector->PointDetectLandmarks(image_data, faces[0], points);
    
    // Visualize the results
    cvRectangle(img_color, cvPoint(faces[0].bbox.x, faces[0].bbox.y), cvPoint(faces[0].bbox.x + faces[0].bbox.width - 1, faces[0].bbox.y + faces[0].bbox.height - 1), CV_RGB(255, 0, 0));
    for (int i = 0; i<pts_num; i++)
    {
        cvCircle(img_color, cvPoint(points[i].x, points[i].y), 2, CV_RGB(0, 255, 0), CV_FILLED);
    }
    //cvSaveImage("result.jpg", img_color);
    cvSaveImage(res_img.c_str(), img_color);
    
    // Release memory
    cvReleaseImage(&img_color);
    cvReleaseImage(&img_grayscale);
    delete[]data;
    return 0;
}

int SeetaAlignment::alignTarget(const IplImage *in_img, string &res_img)
{
    return 0;
}
